import os.path

from torch import nn,optim
import torch
from data import *
from net import *
from torch.utils.data import DataLoader
from torchvision.utils import save_image
device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
weith_path='params/wnet.pth'
data_path=r'D:\datatrain\VOCdevkit\VOC2007'
save_path='train_image'

if __name__ == '__main__':
    dataloader=DataLoader(Mydataset(data_path),batch_size=1,shuffle=True)
    net=Unet().to(device)
    if os.path.exists(weith_path):
        net.load_state_dict(torch.load(weith_path))
        print("load  success")
    else:
        print("not success")

    opt=optim.Adam(net.parameters())
    loss_fun=nn.BCELoss()

    flag=1
    while True:
        for i, (image, segment_image) in enumerate(dataloader):
            image, segment_image = image.to(device), segment_image.to(device)
            out_image=net(image)

            train_loss=loss_fun(out_image,segment_image)
            opt.zero_grad()
            train_loss.backward()
            opt.step()

            if i%5==0:
                print(f'{flag}-{i}-Train_loss=====>{train_loss.item()}')


            if i%50==0:
                torch.save(net.state_dict(),weith_path)

            _image=image[0]
            _segment_image=segment_image[0]
            _out_image=out_image[0]

            img=torch.stack([_image,_segment_image,_out_image],dim=0)
            save_image(img,f'{save_path}/{i}.png')
        flag+=1

